package com.se.net;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.date.TimeInterval;
import cn.hutool.core.net.Ipv4Util;
import com.sun.jna.Memory;
import com.sun.jna.platform.win32.*;
import com.sun.jna.ptr.IntByReference;
import com.util.BitUtil;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

/**
 * @author hdu_huang
 * @since 2024/3/21 11:30
 * 管理信息库(MIB)  Psapi.h IPHlpAPI.h
 *  wmic/netsh     wmic process get name,executablepath,processid
 */
@Slf4j
public class ProcessUtil {

    ////Psapi.GetProcessImageFileName  PsapiUtil  Kernel32Util.QueryFullProcessImageName
    public static void main(String[] args) {

        int pid = 12788;
        String fullProcessName = getFullProcessName(pid);
        String processName = getProcessName(pid);

        System.out.println(fullProcessName);
        System.out.println(processName);


        TimeInterval timer = DateUtil.timer();

        //将内存中的C结构体转换为Java对象
        IPHlpAPI.MIB_TCPTABLE_OWNER_PID tcpTable = getMIB_TCPTABLE_OWNER_PID();

        System.out.println("useTime：" + timer.intervalRestart());
        System.out.println("dwNumEntries:" + tcpTable.dwNumEntries);





        //IPHlpAPI.MIB_TCP_STATE    都是小端表示
        for (int i = 0; i < tcpTable.dwNumEntries; i++) {
            final IPHlpAPI.MIB_TCPROW_OWNER_PID row = tcpTable.table[i];

            if (row.dwOwningPid == 0) {
                continue;
            }


            System.out.println(new ProcessTCPInfo(row));
        }

        System.out.println("useTime：" + timer.interval());



    }

    @Getter
    static enum TCPState {

        CLOSED(1, "Closed"),
        LISTEN(2, "Listen"),
        SYN_SENT(3, "Syn Sent"),
        SYN_RECEIVED(4, "Syn Received"),
        ESTABLISHED(5, "Established"),
        FIN_WAIT_1(6, "Fin Wait 1"),
        FIN_WAIT_2(7, "Fin Wait 2"),
        CLOSE_WAIT(8, "Close Wait"),  //被动关闭的一方, 收到fin后, 还没发fin, 持续时间应该很短
        CLOSING(9, "Closing"),
        LAST_ACK(10, "Last Ack"),
        TIME_WAIT(11, "Time Wait"),   //major 主动关闭的一方, 发送最后一个ack后等待2MSL
        DELETE_TCB(12, "Delete Tcb");

        private final int state;
        private final String message;

        TCPState(int state, String name) {
            this.state = state;
            this.message = name;
        }

        public static TCPState valueOf(Integer value) {

            if (value == null) {
                return null;
            }

            for (TCPState type : TCPState.values()) {
                if (value.equals(type.state)) {
                    return type;
                }
                continue;
            }
            return null;//错误的value
        }
    }

    @Data
    static class ProcessTCPInfo {

        int pid;
        int localPort;
        int remotePort;
        String localIp;
        String remoteIp;
        int state;
        String processName;
        String strState;

        public ProcessTCPInfo(int pid, int localPort, int remotePort, String localIp, String remoteIp, int state) {
            this.pid = pid;
            this.localPort = localPort;
            this.remotePort = remotePort;
            this.localIp = localIp;
            this.remoteIp = remoteIp;
            this.state = state;
        }

        public ProcessTCPInfo(IPHlpAPI.MIB_TCPROW_OWNER_PID tcpRowOwner) {
            transfer(tcpRowOwner);

            strState = TCPState.valueOf(state).getMessage();
        }

        private void transfer(IPHlpAPI.MIB_TCPROW_OWNER_PID tcpRowOwner) {

            this.pid = tcpRowOwner.dwOwningPid;
            this.localPort = BitUtil.littleEndianShortToInt((short) tcpRowOwner.dwLocalPort);
            this.remotePort = BitUtil.littleEndianShortToInt((short) tcpRowOwner.dwRemotePort);
            this.localIp = BitUtil.int2ip(tcpRowOwner.dwLocalAddr, false);
            this.remoteIp = BitUtil.int2ip(tcpRowOwner.dwRemoteAddr, false);

            this.state = tcpRowOwner.dwState;
            this.processName = ProcessUtil.getProcessName(this.pid);

        }

        public String toString() {
            String src = localIp + ":" + localPort;
            String dst = remoteIp + ":" + remotePort;
            return "ProcessUtil.ProcessTCPInfo(pid=" + this.getPid() + ", src=" + src + ", dst=" + dst + ", state=" + strState + ", processName=" + this.getProcessName() + ")";
        }

    }

    //get

    private static IPHlpAPI.MIB_TCPTABLE_OWNER_PID getMIB_TCPTABLE_OWNER_PID() {

        IPHlpAPI instance = IPHlpAPI.INSTANCE;
        //预估计缓冲区大小16k
        IntByReference size = new IntByReference(1024 * 16);
        Memory buffer = new Memory(size.getValue());

        final int TCP_Table_Owner_Pid_All = IPHlpAPI.TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL;

        //返回值: NO_ERROR   ERROR_INSUFFICIENT_BUFFER[pdwSize太小]  ERROR_INVALID_PARAMETER
        //params: 1.返回查询结构体指针 2.预估大小(字节) 如果此值设置得太小，此函数将返回 ERROR_INSUFFICIENT_BUFFER，并且此字段将包含结构的正确大小。
        //3.是否需要排序 4.AF_INET还是AF_INET6   5.表示结构体的种类   6.保留此值必须为零。
        int result = instance.INSTANCE.GetExtendedTcpTable(buffer, size, true, IPHlpAPI.AF_INET, TCP_Table_Owner_Pid_All, 0);

        // 缓冲区大小合适
        if (result == WinError.NO_ERROR) {
            return new IPHlpAPI.MIB_TCPTABLE_OWNER_PID(buffer);
        }

        // 缓冲区大小不够
        if (result != WinError.ERROR_INSUFFICIENT_BUFFER) {
            return null;
        }

        // 分配足够的内存空间
        buffer = new Memory(size.getValue());

        // 第二次调用获取TCP表数据
        result = instance.GetExtendedTcpTable(buffer, size, true, IPHlpAPI.AF_INET, TCP_Table_Owner_Pid_All, 0);

        if (result != WinError.NO_ERROR) {
            return null;
        }

        //将内存中的C结构体转换为Java对象
        return new IPHlpAPI.MIB_TCPTABLE_OWNER_PID(buffer);
    }

    /**
     * 查询给定进程ID的完整可执行文件名称。
     *
     * @param pid  进程ID。
     * @return 完整的进程名称, 如果不存在返回null。 eg:C:\Windows\System32\cmd.exe
     */
    public static String getFullProcessName(int pid) {

        //忽略系统进程
        //pid:0 System Idle Process
        //pid:4 System Process
        if (pid <= 10) return null;

        String fullProcessName = null;

        try {
            return QueryFullProcessImageName(pid, 0);

        } catch (Exception e) {
            log.error("pid: {}", pid, e);
        }
        return null;
    }

    /**
     * 查询给定进程ID的可执行文件名称。
     *
     * @param pid  进程ID。
     * @return 进程名称, 如果不存在返回null。 eg:cmd.exe
     */
    public static String getProcessName(int pid) {

        String fullProcessName = getFullProcessName(pid);

        if (fullProcessName == null) {
            return null;
        }

        String[] split = fullProcessName.split("\\\\");
        return split[split.length - 1];
    }

    /**
     * 查询给定进程ID的完整可执行文件名称。
     *
     * @param pid     进程ID。
     * @param dwFlags 用于控制查询信息的标志。0名称应使用 Win32 路径格式,  1.该名称应使用本机系统路径格式。
     * @return 完整的进程名称。
     * @throws Win32Exception 如果打开进程或查询进程图像名称失败，则抛出Win32Exception。
     */
    private static String QueryFullProcessImageName(int pid, int dwFlags) {

        WinNT.HANDLE process = null;
        Win32Exception we = null;

        try {
            process = Kernel32.INSTANCE.OpenProcess(WinNT.PROCESS_QUERY_LIMITED_INFORMATION, false, pid);
            if (process == null) {
                throw new Win32Exception(Kernel32.INSTANCE.GetLastError());
            }
            //String processImagePath = PsapiUtil.GetProcessImageFileName(process);
            return Kernel32Util.QueryFullProcessImageName(process, dwFlags);

        } catch (Win32Exception e) {
            we = e;
            throw we; // re-throw to invoke finally block
        } finally {
            try {
                Kernel32.INSTANCE.CloseHandle(process);
            } catch (Win32Exception e) {
                if (we == null) {
                    we = e;
                } else {
                    we.addSuppressed(e);
                }
            }
            if (we != null) {
                throw we;
            }
        }
    }
}
